import copy
import os
import numpy as np
from astropy.io import fits
from astropy.modeling import fitting
from astropy.modeling.polynomial import Polynomial2D
from astropy.table import Table
from astropy.visualization import ZScaleInterval
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.cm import get_cmap


def apply_gaintab(image, gain_tab, imgconf, return_gain=False):
    data = image.astype(float)
    header = dict()
    for chan in range(imgconf.nchan):
        index = gain_tab['channel'] == chan + 1
        if index.sum() == 1:
            gain = gain_tab['gain'][index][0]
            header['GAIN_{:02d}'.format(chan + 1)] = gain
        else:
            gain = np.nan
        x1 = imgconf.mos_x0[chan]
        x2 = x1 + imgconf.nx[chan]
        y1 = imgconf.mos_y0[chan]
        y2 = y1 + imgconf.ny[chan]
        data[y1:y2, x1:x2] = data[y1:y2, x1:x2] * gain
    if return_gain:
        return data, header
    else:
        return data


def make_gainmap(fgain_tab, fout, imgconf):
    """
    @2024-12-02:将gainmap从64位浮点数改为32位浮点数
    """
    data = np.ones(imgconf.mos_shape, dtype=np.float32)
    tab = Table.read(fgain_tab, format='ipac')
    data, header = apply_gaintab(data, tab, imgconf, return_gain=True)
    if np.isfinite(data).sum() == 0:
        print('Warning: no valid channel to make gain map')
        return False
#    hdu = fits.PrimaryHDU(data)
    hdu = fits.PrimaryHDU(data.astype(np.float32))
    for key in header:
        hdu.header[key] = header[key]
    hdu.writeto(fout, overwrite=True)


def apply_gain(fimage, fgain, fout, ext=None, header_info=None):
    with fits.open(fimage) as img, fits.open(fgain) as gain:
        if ext is None:
            exts = np.arange(len(img))
        else:
            exts = [ext, ]
        hdulist = fits.HDUList([])
        for i in range(len(exts)):
            index = exts[i]
            data = img[index].data * gain[0].data
            header = img[index].header
            if i == 0:
                hdu = fits.PrimaryHDU(data, header=header)
            else:
                hdu = fits.ImageHDU(data, header=header)
            if 'BUNIT' in hdu.header:
                hdu.header['BUNIT'] = hdu.header['BUNIT'].replace('adu', 'ph')
                hdu.header['GAINFILE'] = fgain
                hdu.header['GAIN'] = np.nanmean(gain[0].data)
            if header_info is not None:
                for key in header_info:
                    hdu.header[key] = header_info[key]
            hdulist.append(hdu)
        hdulist.writeto(fout, overwrite=True)


def plot_image(fname, ffig=None, title=None):
    if not os.path.exists(fname):
        return
    if ffig is None:
        ffig = fname.replace('.fits', '.png')
    matplotlib.rc('image', cmap='Blues_r')
    colormap = copy.copy(get_cmap("Blues_r"))
    colormap.set_bad(color='white')
    zs = ZScaleInterval()

    with fits.open(fname) as f:
        plt.figure(figsize=(12, 9))
        index = np.isfinite(f[0].data)
        dd = f[0].data[index]
        lims = zs.get_limits(dd[dd != 0])
        plt.imshow(f[0].data, origin='lower',
                   vmin=lims[0], vmax=lims[1], cmap=colormap)
        if 'BUNIT' in f[0].header:
            unit = f[0].header['BUNIT'].replace('adu', 'ADU').replace('ph', 'e-')
        else:
            unit = 'ADU'
        plt.colorbar(label=unit, extendfrac='auto')

        if title is None:
            title = f[0].header['OBSTYPE'].strip()
        if title == 'flat':
            title = 'flat: {:.1f} nm'.format(f[0].header['WAVELEN'])
        elif title == 'bias' or title == 'dark':
            temp = f[0].header['SYSTEMP']
            if temp < 998:
                title = '{}: {:.1f} {}C'.format(title, temp - 273, u"\u00b0")
        else:
            pass
        plt.title(title)
        plt.tight_layout()
        plt.savefig(ffig, dpi=100)
        plt.close()


def plot_tsys(flist, ffig, n_col=3, fgain=None):
    zs = ZScaleInterval()
    colormap = 'Blues_r'

    gainmap = None
    if fgain is not None:
        gainmap = fits.getdata(fgain)
    n = len(flist)
    if n <= 1:
        print('too few data, skip')
        return False
    flist.sort(key=lambda x: fits.getheader(x)['SYSTEMP'])

    # plot the image
    n_row = n // n_col
    if n - n_row * n_col > 0:
        n_row += 1
    plt.figure(figsize=(18, 5 * n_row))
    for i in range(n):
        plt.subplot(n_row, 3, i+1)
        with fits.open(flist[i]) as img:
            data = img[0].data
            if fgain is not None:
                data *= gainmap
            if 'BUNIT' in img[0].header:
                unit = img[0].header['BUNIT'].replace('ph', 'e-').replace('adu', 'ADU')
            index = data > 0
            vmin, vmax = zs.get_limits(data[index])
            plt.imshow(data, vmin=vmin, vmax=vmax, cmap=colormap, origin='lower')
            plt.title('{:.1f} degree, {:.1f} s'.format(img[0].header['SYSTEMP'] - 273.15, img[0].header['DARKTIME']))
            plt.yticks([])
            plt.xticks([])
            plt.colorbar(label=unit)

    # output
    plt.tight_layout()
    plt.savefig(ffig)
    plt.close()


def plot_wave(flist, ffig, n_col=2, fgain=None):
    n = len(flist)
    if n <= 1:
        return False
    wave = [fits.getheader(f)['WAVELEN'] for f in flist]
    index = np.argsort(np.asarray(wave))
    namelist = np.asarray(flist)[index]
    zs = ZScaleInterval()
    gain = None
    if fgain is not None:
        gain = fits.getdata(fgain)

    # plot the image
    n_row = n // n_col
    if n - n_row * n_col > 0:
        n_row += 1
    plt.figure(figsize=(5 * n_col, 4 * n_row))
    for i in range(n):
        plt.subplot(n_row, n_col, i+1)
        with fits.open(namelist[i]) as img:
            data = img[0].data
            if fgain is not None:
                data = data * gain
            index = np.isfinite(data)
            data /= data[index].mean()
#            vmin, vmax = zs.get_limits(data[index])
            vmin, vmax = 0.9, 1.05 # added @20241116
            
            #############################
            # remove starnge lines
            # added@20241116
            rm = np.median(data)
            rms = np.std(data)
            idx1 = data < rm - 5*rms
            idx2 = data > rm + 5*rms
            data[idx1] = rm
            data[idx2] = rm
            #############################
            
            plt.imshow(data, vmin=vmin, vmax=vmax, cmap='Blues_r', origin='lower')
            plt.title('wavelength = {:.1f} nm'.format(img[0].header['WAVELEN']))
            plt.yticks([])
            plt.xticks([])
            plt.colorbar()

    # output
    plt.tight_layout()
    plt.savefig(ffig)
    plt.close()


def get_dark(fdark, fbias, fout):
    with fits.open(fdark) as img, fits.open(fbias) as bias:
        data = (img[0].data - bias[0].data) / img[0].header['DARKTIME']
        hdu = fits.PrimaryHDU(data, header=img[0].header)
        hdu.header['BUNIT'] = 'adu/s'
        hdu.writeto(fout, overwrite=True)


def polyfit(data, degree=3):
    ny, nx = data.shape
    yy, xx = np.mgrid[:ny, :nx]
    p_init = Polynomial2D(degree=degree)
    fit_p = fitting.LinearLSQFitter()
    p = fit_p(p_init, xx, yy, data)
    model = p(xx, yy)
    return model


def make_cutout(image, xcent, ycent, xsize, ysize, return_corner=False):
    ny, nx = image.shape
    x1 = max(xcent - xsize // 2, 0)
    x2 = min(xcent + xsize // 2, nx)
    y1 = max(ycent - ysize // 2, 0)
    y2 = min(ycent + ysize // 2, ny)
    output = image[y1:y2, x1:x2]
    if return_corner:
        return output, x1, x2, y1, y2
    else:
        return output


def subtract_bias(fimage, fbias, fout):
    with fits.open(fimage) as img, fits.open(fbias) as bias:
        data = img[0].data - bias[0].data
        hdu = fits.PrimaryHDU(data, header=img[0].header)
        if 'DARKTIME' in img[0].header and 'DARKTIME' in bias[0].header:
            hdu.header['DARKTIME'] = img[0].header['DARKTIME'] - bias[0].header['DARKTIME']
        hdu.writeto(fout, overwrite=True)
